import csv
import numpy as np
from scipy.interpolate import RegularGridInterpolator
from scipy.spatial import KDTree
from numba import njit
# import time


class NodeInterpolator:
    def __init__(self, grid_data_path, result_coordinates_path, output_file_path="Grid_new_and_data.txt",
                 is_regular=True, k=4):
        """
        初始化插值器类
        参数:
            grid_data_path (str): 网格和数据文件路径
            result_coordinates_path (str): 结果网格坐标文件路径
            output_file_path (str): 插值结果保存路径，默认为 'Grid_new_and_data.txt'
            is_regular  # 默认为True，代表原坐标为规则网格；若为False,则代表不确定，由_is_regular_grid函数进行判断。
            k (int): 非规则网格时用于插值的最近邻点数量，默认为 4
        """
        self.grid_data_path = grid_data_path
        self.result_coordinates_path = result_coordinates_path
        self.output_file_path = output_file_path
        self.is_regular = is_regular  # 默认为True
        self.k = k  # 非规则网格的最近邻点数量
        self.grid_data = self._read_file(self.grid_data_path)
        self.result_coordinates = self._read_file(self.result_coordinates_path)
        self.x, self.y, self.z = self.unique_no_sort(self.grid_data[:, 0]), self.unique_no_sort(
            self.grid_data[:, 1]), self.unique_no_sort(self.grid_data[:, 2])  # 获取唯一的XYZ坐标

    def _read_file(self, file_path):
        """
        读取数据文件
        参数:
            file_path (str): 文件路径
        返回:
            np.ndarray: 读取的文件数据
        """
        return np.loadtxt(file_path, skiprows=1, delimiter='\t', encoding='utf-16le')

    def _is_regular_grid(self, x_unique, y_unique, z_unique, tolerance=0.5):
        """
        检查 x, y, z 是否分别等间距，允许误差在指定容差内。
        等间距的定义是数组元素间距的最大值和最小值之差不超过容差。
        参数：
            x_unique (array-like): x 方向唯一值的数组。
            y_unique (array-like): y 方向唯一值的数组。
            z_unique (array-like): z 方向唯一值的数组。
            tolerance (float): 允许的最大间距变化范围,具有默认值。
        返回：
            bool: 如果 x, y, z 都等间距返回 True，否则返回 False。
        """

        def is_equally_spaced(arr):
            """
            检查单个数组是否等间距。
            """
            if len(arr) < 2:
                return True
            diffs = np.abs(np.diff(arr))
            return np.max(diffs) - np.min(diffs) <= tolerance

        return (
                is_equally_spaced(x_unique) and
                is_equally_spaced(y_unique) and
                is_equally_spaced(z_unique)
        )

    def unique_no_sort(self, array):
        # 获取唯一值和它们在原始数组中的索引   按照原始数组的顺序排序唯一值
        array = np.round(array, decimals=4)
        unique_values, indices = np.unique(array, return_index=True)
        unique_values_in_order = array[np.sort(indices)]
        return unique_values_in_order

    @staticmethod
    @njit
    def _weighted_interpolation(distances, values):
        """
        使用距离加权的插值方法
        参数:
            distances (np.ndarray): 距离数组
            values (np.ndarray): 值数组
        返回:
            float: 插值结果
        """
        weights = 1 / (distances + 1e-12)  # 避免除零
        weights /= weights.sum()  # 归一化权重
        return np.dot(weights, values)

    def _interpolate_with_kdtree(self, points, values, query_points):
        """
        使用 KDTree 和加权插值方法进行非规则网格插值
        参数:
            points (np.ndarray): 网格点坐标
            values (np.ndarray): 网格点的值
            query_points (np.ndarray): 查询点坐标
        返回:
            np.ndarray: 插值结果
        """
        tree = KDTree(points)

        distances, indices = tree.query(query_points, k=self.k)
        interpolated_values = np.zeros(len(query_points))

        for i in range(len(query_points)):
            interpolated_values[i] = self._weighted_interpolation(distances[i], values[indices[i]])

        return interpolated_values

    def perform_interpolation(self):
        """
        执行插值并保存结果
        """
        if self.is_regular:
            # 规则
            s1_values = self.grid_data[:, 3].reshape(len(self.y), len(self.x), len(self.z), order='F')
            s1_values = np.transpose(s1_values, (1, 0, 2))  # 确保最终顺序为 (x, y, z)
            interpolator = RegularGridInterpolator((self.x, self.y, self.z), s1_values, method='linear')  # 哪个好？！！
            interpolated_values = interpolator(self.result_coordinates[:, 1:4])
        elif self._is_regular_grid(self.x, self.y, self.z):
            # 规则
            s1_values = self.grid_data[:, 3].reshape(len(self.y), len(self.x), len(self.z), order='F')
            s1_values = np.transpose(s1_values, (1, 0, 2))  # 确保最终顺序为 (x, y, z)
            interpolator = RegularGridInterpolator((self.x, self.y, self.z), s1_values, method='linear')
            interpolated_values = interpolator(self.result_coordinates[:, 1:4])
        else:
            # 非规则
            points, values = self.grid_data[:, :3], self.grid_data[:, 3]
            interpolated_values = self._interpolate_with_kdtree(points, values, self.result_coordinates[:, 1:4])
        # 保存插值结果
        self._save_results(interpolated_values)

    def _save_results(self, interpolated_values):
        """
        保存插值结果到文件
        参数:
            interpolated_values (np.ndarray): 插值结果数组
        """
        data_to_save = np.column_stack((self.result_coordinates[:, 1:4], interpolated_values))
        with open(self.output_file_path, 'w', encoding='utf-16le', newline='') as file:
            writer = csv.writer(file, delimiter='\t')
            writer.writerow(['x', 'y', 'z', 'S1'])
            writer.writerows(data_to_save)


# start = time.time()
interpolator = NodeInterpolator("../Question2/Grid_and_data.txt",
                                "../Question2/Result_grid_coordinates.txt",
                                "../Question2/Grid_D_and_data.txt")
interpolator.perform_interpolation()
# print(time.time() - start)
